#!/bin/bash

# 设置颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# 打印带颜色的信息
print_info() {
    echo -e "${GREEN}[INFO] $1${NC}"
}

print_warn() {
    echo -e "${YELLOW}[WARN] $1${NC}"
}

print_error() {
    echo -e "${RED}[ERROR] $1${NC}"
}

# 检查是否为root用户
check_root() {
    if [ "$EUID" -ne 0 ]; then
        print_error "请使用root权限运行此脚本"
        exit 1
    fi
}

# 更新系统
update_system() {
    print_info "更新系统包..."
    apt update && apt upgrade -y
}

# 安装基础工具
install_basic_tools() {
    print_info "安装基础工具..."
    apt install -y vim wget curl git net-tools htop iotop iftop ntp ntpdate
}

# 配置主机名和hosts
configure_hosts() {
    print_info "配置hosts文件..."
    cat >> /etc/hosts << EOF
# 主节点（Master Node）
10.24.8.13 hadoop-master

# 工作节点（Worker Nodes）
10.24.8.133 hadoop-worker1
10.24.8.96 hadoop-worker2
EOF

    # 根据IP地址设置当前主机名
    local current_ip=$(hostname -I | awk '{print $1}')
    case $current_ip in
        "10.24.8.13")
            hostnamectl set-hostname hadoop-master
            ;;
        "10.24.8.133")
            hostnamectl set-hostname hadoop-worker1
            ;;
        "10.24.8.96")
            hostnamectl set-hostname hadoop-worker2
            ;;
    esac
}

# 配置时间同步
configure_ntp() {
    print_info "配置NTP时间同步..."
    systemctl enable ntp
    systemctl start ntp
    ntpdate cn.pool.ntp.org
}

# 配置系统参数
configure_system_params() {
    print_info "配置系统参数..."
    
    # 配置系统限制
    cat >> /etc/security/limits.conf << EOF
* soft nofile 65535
* hard nofile 65535
* soft nproc 65535
* hard nproc 65535
EOF

    # 配置内核参数
    cat >> /etc/sysctl.conf << EOF
# 网络参数
net.ipv4.tcp_syncookies = 1
net.ipv4.tcp_tw_reuse = 1
net.ipv4.tcp_tw_recycle = 1
net.ipv4.tcp_fin_timeout = 30
net.ipv4.tcp_keepalive_time = 1200
net.ipv4.ip_local_port_range = 1024 65535

# 内存参数
vm.swappiness = 10
vm.dirty_ratio = 60
vm.dirty_background_ratio = 2

# 文件系统参数
fs.file-max = 65535
EOF

    # 应用内核参数
    sysctl -p
}

# 创建Hadoop用户和组
create_hadoop_user() {
    print_info "创建Hadoop用户和组..."
    groupadd hadoop
    useradd -g hadoop -m hadoop
    echo "hadoop:hadoop" | chpasswd
    
    # 配置sudo权限
    echo "hadoop ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers
}

# 创建必要的目录
create_directories() {
    print_info "创建必要的目录..."
    mkdir -p /opt/hadoop
    mkdir -p /opt/hadoop/data
    mkdir -p /opt/hadoop/logs
    mkdir -p /opt/hadoop/tmp
    
    chown -R hadoop:hadoop /opt/hadoop
}

# 配置SSH免密登录
configure_ssh() {
    print_info "配置SSH免密登录..."
    su - hadoop -c "ssh-keygen -t rsa -P '' -f ~/.ssh/id_rsa"
    
    # 将公钥复制到其他节点（需要手动执行）
    print_warn "请手动将SSH公钥复制到其他节点："
    print_warn "ssh-copy-id hadoop@hadoop-master"
    print_warn "ssh-copy-id hadoop@hadoop-worker1"
    print_warn "ssh-copy-id hadoop@hadoop-worker2"
}

# 主函数
main() {
    print_info "开始环境准备..."
    check_root
    update_system
    install_basic_tools
    configure_hosts
    configure_ntp
    configure_system_params
    create_hadoop_user
    create_directories
    configure_ssh
    print_info "环境准备完成！"
}

# 执行主函数
main 